"""
MySQL数据库工具类
提供MySQL数据库连接、指定schema、增删改查操作以及存储过程支持
"""

from typing import Any, List, Dict, Optional, Union


class MySQLUtils:
    """
    MySQL数据库操作工具类
    支持连接管理、指定schema、增删改查操作、存储过程调用和占位符操作
    
    注意：此类需要mysql-connector-python或PyMySQL库支持
    安装方式：
    pip install mysql-connector-python
    或
    pip install PyMySQL
    """

    def __init__(self, host: str, port: int, username: str, password: str, 
                 database: str = None, schema: str = None, charset: str = 'utf8'):
        """
        初始化MySQL连接参数
        
        Args:
            host: 数据库主机地址
            port: 数据库端口
            username: 用户名
            password: 密码
            database: 数据库名（可选）
            schema: schema名（可选）
            charset: 字符集，默认为utf8
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.database = database
        self.schema = schema
        self.charset = charset
        self.connection = None
        self.db_type = "mysql"

    def connect(self) -> 'MySQLUtils':
        """
        建立数据库连接
        
        Returns:
            MySQLUtils实例，支持链式调用
        """
        try:
            # 尝试使用mysql-connector-python
            try:
                import mysql.connector
                self.connection = mysql.connector.connect(
                    host=self.host,
                    port=self.port,
                    user=self.username,
                    password=self.password,
                    database=self.database,
                    charset=self.charset,
                    autocommit=True
                )
                self.db_type = "mysql-connector"
                
                # 如果指定了schema，则使用该schema
                if self.schema:
                    with self.connection.cursor() as cursor:
                        cursor.execute(f"USE {self.schema}")
                        
            except ImportError:
                # 如果mysql-connector-python不可用，尝试使用PyMySQL
                try:
                    import pymysql
                    self.connection = pymysql.connect(
                        host=self.host,
                        port=self.port,
                        user=self.username,
                        password=self.password,
                        database=self.database,
                        charset=self.charset,
                        autocommit=True
                    )
                    self.db_type = "pymysql"
                    
                    # 如果指定了schema，则使用该schema
                    if self.schema:
                        with self.connection.cursor() as cursor:
                            cursor.execute(f"USE {self.schema}")
                            
                except ImportError:
                    raise ImportError("需要安装mysql-connector-python或PyMySQL库。请运行: pip install mysql-connector-python 或 pip install PyMySQL")
                    
            return self
        except Exception as e:
            raise Exception(f"连接MySQL数据库失败: {str(e)}")

    def disconnect(self) -> None:
        """
        关闭数据库连接
        """
        if self.connection:
            if hasattr(self.connection, 'is_connected'):
                # mysql-connector-python
                if self.connection.is_connected():
                    self.connection.close()
            else:
                # PyMySQL
                self.connection.close()
            self.connection = None

    def execute_query(self, sql: str, params: Optional[Union[tuple, dict]] = None) -> List[Dict]:
        """
        执行查询语句
        
        Args:
            sql: SQL查询语句，支持占位符
            params: 查询参数，可以是元组或字典
            
        Returns:
            查询结果列表
        """
        if not self.connection:
            raise Exception("数据库未连接")
            
        try:
            # 判断使用的数据库驱动类型
            if hasattr(self.connection, 'cursor'):
                if self.db_type == "mysql-connector" and hasattr(self.connection, 'is_connected') and self.connection.is_connected():
                    # mysql-connector-python
                    cursor = self.connection.cursor(dictionary=True)
                elif self.db_type == "pymysql":
                    # PyMySQL
                    import pymysql.cursors
                    cursor = self.connection.cursor(pymysql.cursors.DictCursor)
                else:
                    raise Exception("数据库连接已断开")
                    
                cursor.execute(sql, params)
                result = cursor.fetchall()
                cursor.close()
                return result
            else:
                raise Exception("无效的数据库连接")
        except Exception as e:
            raise Exception(f"执行查询失败: {str(e)}")

    def execute_update(self, sql: str, params: Optional[Union[tuple, dict]] = None) -> int:
        """
        执行更新语句（INSERT, UPDATE, DELETE）
        
        Args:
            sql: SQL更新语句，支持占位符
            params: 更新参数，可以是元组或字典
            
        Returns:
            受影响的行数
        """
        if not self.connection:
            raise Exception("数据库未连接")
            
        try:
            # 判断使用的数据库驱动类型
            if hasattr(self.connection, 'cursor'):
                if self.db_type == "mysql-connector" and hasattr(self.connection, 'is_connected') and self.connection.is_connected():
                    # mysql-connector-python
                    cursor = self.connection.cursor()
                elif self.db_type == "pymysql":
                    # PyMySQL
                    cursor = self.connection.cursor()
                else:
                    raise Exception("数据库连接已断开")
                    
                affected_rows = cursor.execute(sql, params)
                if hasattr(self.connection, 'commit'):
                    self.connection.commit()
                cursor.close()
                return affected_rows
            else:
                raise Exception("无效的数据库连接")
        except Exception as e:
            if hasattr(self.connection, 'rollback'):
                self.connection.rollback()
            raise Exception(f"执行更新失败: {str(e)}")

    def execute_procedure(self, procedure_name: str, 
                         params: Optional[Union[tuple, dict]] = None) -> Any:
        """
        执行存储过程
        
        Args:
            procedure_name: 存储过程名称
            params: 存储过程参数，可以是元组或字典
            
        Returns:
            存储过程执行结果
        """
        if not self.connection:
            raise Exception("数据库未连接")
            
        try:
            # 判断使用的数据库驱动类型
            if hasattr(self.connection, 'cursor'):
                if self.db_type == "mysql-connector" and hasattr(self.connection, 'is_connected') and self.connection.is_connected():
                    # mysql-connector-python
                    cursor = self.connection.cursor(dictionary=True)
                elif self.db_type == "pymysql":
                    # PyMySQL
                    import pymysql.cursors
                    cursor = self.connection.cursor(pymysql.cursors.DictCursor)
                else:
                    raise Exception("数据库连接已断开")
                
                # 构造调用存储过程的SQL
                if params:
                    if isinstance(params, dict):
                        # 处理命名参数
                        param_names = ', '.join([f"@{k} := %({k})s" for k in params.keys()])
                        sql = f"CALL {procedure_name}({param_names})"
                        cursor.execute(sql, params)
                    else:
                        # 处理位置参数
                        placeholders = ', '.join(['%s' for _ in params])
                        sql = f"CALL {procedure_name}({placeholders})"
                        cursor.execute(sql, params)
                else:
                    sql = f"CALL {procedure_name}()"
                    cursor.execute(sql)
                    
                result = cursor.fetchall()
                cursor.close()
                return result
            else:
                raise Exception("无效的数据库连接")
        except Exception as e:
            raise Exception(f"执行存储过程失败: {str(e)}")

    def insert(self, table: str, data: Dict[str, Any]) -> int:
        """
        插入数据
        
        Args:
            table: 表名
            data: 要插入的数据字典
            
        Returns:
            受影响的行数
        """
        if not data:
            raise ValueError("数据不能为空")
            
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s' for _ in data])
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        
        return self.execute_update(sql, tuple(data.values()))

    def update(self, table: str, data: Dict[str, Any], condition: str, 
               condition_params: Optional[tuple] = None) -> int:
        """
        更新数据
        
        Args:
            table: 表名
            data: 要更新的数据字典
            condition: 更新条件
            condition_params: 条件参数
            
        Returns:
            受影响的行数
        """
        if not data:
            raise ValueError("更新数据不能为空")
            
        set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
        sql = f"UPDATE {table} SET {set_clause} WHERE {condition}"
        
        params = list(data.values())
        if condition_params:
            params.extend(condition_params)
            
        return self.execute_update(sql, tuple(params))

    def delete(self, table: str, condition: str, 
               condition_params: Optional[tuple] = None) -> int:
        """
        删除数据
        
        Args:
            table: 表名
            condition: 删除条件
            condition_params: 条件参数
            
        Returns:
            受影响的行数
        """
        sql = f"DELETE FROM {table} WHERE {condition}"
        return self.execute_update(sql, condition_params)

    def select(self, table: str, columns: str = "*", condition: str = "1=1", 
               condition_params: Optional[tuple] = None) -> List[Dict]:
        """
        查询数据
        
        Args:
            table: 表名
            columns: 要查询的列，默认为所有列(*)
            condition: 查询条件，默认为1=1（查询所有）
            condition_params: 条件参数
            
        Returns:
            查询结果列表
        """
        sql = f"SELECT {columns} FROM {table} WHERE {condition}"
        return self.execute_query(sql, condition_params)

    def select_one(self, table: str, columns: str = "*", condition: str = "1=1", 
                   condition_params: Optional[tuple] = None) -> Optional[Dict]:
        """
        查询单条数据
        
        Args:
            table: 表名
            columns: 要查询的列，默认为所有列(*)
            condition: 查询条件，默认为1=1（查询所有）
            condition_params: 条件参数
            
        Returns:
            单条查询结果，如果没有则返回None
        """
        result = self.select(table, columns, condition, condition_params)
        return result[0] if result else None

    def __enter__(self) -> 'MySQLUtils':
        """
        上下文管理器入口，自动连接数据库
        """
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """
        上下文管理器出口，自动关闭数据库连接
        """
        self.disconnect()